/*
*    Phantom v2.0 GUI Automation Software
*    Copyright (C) 2009  John C. Van Arsdall (john@phantomtest.com)
*
*    This program is free software; you can redistribute it and/or modify
*    it under the terms of the GNU General Public License as published by
*    the Free Software Foundation; either version 2 of the License, or
*    (at your option) any later version.
*
*    This program is distributed in the hope that it will be useful,
*    but WITHOUT ANY WARRANTY; without even the implied warranty of
*    MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
*    GNU General Public License for more details.
*
*    You should have received a copy of the GNU General Public License along
*    with this program; if not, write to the Free Software Foundation, Inc.,
*    51 Franklin Street, Fifth Floor, Boston, MA 02110-1301 USA.
*
*/
#include ".\ssi_socket.h"

#include <winsock.h>
#include <stdlib.h>
#include <stdio.h>
#include <time.h>

int CSocketSSI::do_log = 1;
char CSocketSSI::socket_log_file[1024] = "";

CSocketSSI::CSocketSSI(void)
: myPort(0)
{
	myHandle = 0;
}

CSocketSSI::~CSocketSSI(void)
{
  // Ensure disconnects are handled explicitely
  //if(myHandle != 0){
  //  CSocketSSI::Disconnect(myHandle);
  //  myHandle = 0;
  //}
}

int CSocketSSI::Init(void)
{
  WSAData wsaData;
  int nCode;
  nCode = WSAStartup(MAKEWORD(1, 1), &wsaData);
  
	return nCode;

}

SSI_ADDR CSocketSSI::LookupAddress(const char* pcHost)
{
	unsigned long nRemoteAddr = inet_addr(pcHost);
	if (nRemoteAddr == INADDR_NONE) {
		// pcHost isn't a dotted IP, so resolve it through DNS
		// Do not attempt to free this memory...
		hostent* pHE = gethostbyname(pcHost);
		if (pHE == 0) {
			return INADDR_NONE;
		}
		nRemoteAddr = *((unsigned long*)pHE->h_addr_list[0]);
	}

  //in_addr Address;
  //memcpy(&Address, &nRemoteAddr, sizeof(u_long)); 

	return nRemoteAddr;
}

SSI_SOCKET CSocketSSI::Connect(SSI_ADDR Address, int Port)
{
  SSI_SOCKET sd = (SSI_SOCKET)socket(AF_INET, SOCK_STREAM, 0);
  if (sd != INVALID_SOCKET) {
    sockaddr_in sinRemote;
    sinRemote.sin_family = AF_INET;
    sinRemote.sin_addr.s_addr = Address;
    sinRemote.sin_port = htons(Port);

		//bool bval = false;
		//setsockopt(sd, SOL_SOCKET, SO_REUSEADDR, (const char*)&bval, sizeof(bval));


		// Why did I ever use this? (20 July 2007)
		//if(bind(sd, (sockaddr*)&sinRemote, sizeof(sockaddr_in)) == SOCKET_ERROR){
		//	DispLastError(__FILE__, __LINE__);
		//	sd = INVALID_SOCKET;
		//	return sd;
		//}

    if (connect(sd, (sockaddr*)&sinRemote, sizeof(sockaddr_in)) == SOCKET_ERROR){
			DispLastError(__FILE__, __LINE__);
			closesocket(sd);
      sd = INVALID_SOCKET;
		}
	}else{
		DispLastError(__FILE__, __LINE__);
	}
  return sd;
}


int CSocketSSI::CleanUp(void)
{
	//return WSACleanup();
  return 1;
}

int CSocketSSI::CleanUpSingle(void)
{
	return WSACleanup();
  //return 1;
}

CSockCmd CSocketSSI::SendCommand(SSI_SOCKET InSock, CSockCmd& InCommand)
{
	CSockCmd Ret;

	int data_size = sizeof(InCommand.ID) + sizeof(InCommand.Size) + InCommand.Size;
	int offset = 0;
	char* data_buf = new char[data_size];
	memcpy(data_buf + offset, &InCommand.ID, sizeof(InCommand.ID));
	offset += sizeof(InCommand.ID);

	memcpy(data_buf + offset, &InCommand.Size, sizeof(InCommand.Size));
	offset += sizeof(InCommand.Size);

	memcpy(data_buf + offset, InCommand.Data, InCommand.Size);

	if(send(InSock, data_buf, data_size, 0) == SOCKET_ERROR){
		DispLastError(__FILE__, __LINE__);
		Ret.ID = RSP_ERROR;
		return Ret;
	}else{
		Ret.ID = RSP_ACK;
	}

	return Ret;
}

CSockCmd CSocketSSI::ReadCommand(SSI_SOCKET InSock)
{
	// Read reply from server
	CSockCmd Ret;
	int cmd_id = 0;
	int data_size = 0;

	int read_bytes = 0;

	while(read_bytes < sizeof(cmd_id)){
		char* buff = (char*)&cmd_id;
		int iRead = 0;
		iRead = recv(InSock, &buff[read_bytes], 1, 0);
		if(iRead <= 0){
			Ret.ID = RSP_ERROR;
			return Ret;
		}

		read_bytes += iRead;
		//printf("%d bytes read for Command: %d\n", read_bytes, cmd_id);
		if(read_bytes > sizeof(cmd_id)){
			Ret.ID = RSP_ERROR;
			return Ret;
		}
	}

	//Log("Socket: Reading ID");
	read_bytes = 0;
	while(read_bytes < sizeof(cmd_id)){
		char* buff = (char*)&data_size;
		int iRead = 0;
		iRead = recv(InSock, &buff[read_bytes], 1, 0);
		if(iRead <= 0){
			Ret.ID = RSP_ERROR;
			return Ret;
		}
		read_bytes += iRead;
		if(read_bytes > sizeof(data_size)){
			Ret.ID = RSP_ERROR;
			return Ret;
		}
	}

	char* ReadBuffer = 0;
	if(data_size > 0){
		ReadBuffer = new char[data_size+1];
		read_bytes = recv(InSock, ReadBuffer, data_size, 0);
		if(read_bytes != data_size){
			Ret.ID = RSP_ERROR;
			delete ReadBuffer;
			return Ret;
		}
		ReadBuffer[data_size] = 0;
		Ret.Set(cmd_id, data_size, ReadBuffer);
		delete ReadBuffer;
	}else{
		Ret.Set(cmd_id, data_size, 0);
	}
	
	return Ret;
}

int CSocketSSI::Disconnect(SSI_SOCKET sd)
{
    // Disallow any further data sends.  This will tell the other side
    // that we want to go away now.  If we skip this step, we don't
    // shut the connection down nicely.

    char S[1024];
    if(sprintf(S, "DISCONNECT: Socket is shutting down: %s (%d)", IPFromSocket(sd), PortFromSocket(sd))){
      SSI_VERBOSE_LOG(S);
    }
    
    if (shutdown(sd, SSI_SOCKET_SEND) == SOCKET_ERROR) {    
			DispLastError(__FILE__, __LINE__);
      //return 0;  // want to continue shut-down
    }

    // Receive any extra data still sitting on the socket.  After all
    // data is received, this call will block until the remote host
    // acknowledges the TCP control packet sent by the shutdown above.
    // Then we'll get a 0 back from recv, signalling that the remote
    // host has closed its side of the connection.
    char ReadBuffer[4096];
    while (1){
			int nNewBytes = recv(sd, ReadBuffer, 4096, 0);
      if (nNewBytes == SOCKET_ERROR){
				break;       // why was I returning or not breaking?
		//		//return 0;  // want to continue shut-down
      }else if (nNewBytes != 0){
            //cerr << endl << "FYI, received " << nNewBytes <<
            //        " unexpected bytes during shutdown." << endl;
      }else{
            // Okay, we're done!
				break;
			}
    }

		return CloseSocket(sd);
    // Close the socket.
    //if (closesocket(sd) == SOCKET_ERROR) {
		//	DispLastError(__FILE__, __LINE__);
    //  return 0;
    //}
}

int CSocketSSI::CloseSocket(SSI_SOCKET sd){
  char S[1024];
  if(sprintf(S, "CLOSESOCKET Closing: %s (%d)", IPFromSocket(sd), PortFromSocket(sd))){
    CSocketSSI::Log(S);
  }

  if (closesocket(sd) == SOCKET_ERROR) {
		DispLastError(__FILE__, __LINE__);
    return 0;
  }

	return 1;
}

SSI_SOCKET CSocketSSI::Listen(SSI_ADDR Address, int Port)
{
  SSI_SOCKET sd = (SSI_SOCKET)socket(AF_INET, SOCK_STREAM, 0);
  if (sd != INVALID_SOCKET) {
    sockaddr_in sinInterface;
    sinInterface.sin_family = AF_INET;
    sinInterface.sin_addr.s_addr = Address;
    sinInterface.sin_port = htons(Port);
    if(bind(sd, (sockaddr*)&sinInterface, sizeof(sockaddr_in)) != SOCKET_ERROR){
      
      listen(sd, 1);
      
      char S[1024];
      if(sprintf(S, "Listening on: %s [Port ID=%d {%d}]", IPFromSocket(sd), Port, PortFromSocket(sd))){
        CSocketSSI::Log(S);
      }

      return sd;
		}else{
			DispLastError(__FILE__, __LINE__);
		}
	}else{
		DispLastError(__FILE__, __LINE__);
	}

  return INVALID_SOCKET;
}

SSI_SOCKET CSocketSSI::AcceptConnection(SSI_SOCKET InSock, int& Port)
{
	sockaddr_in sinRemote;
	int nAddrSize = sizeof(sinRemote);
	SSI_SOCKET Ret = (SSI_SOCKET)accept(InSock, (sockaddr*)&sinRemote, &nAddrSize);
	//myAddr = sinRemote;
	int aPort = (int)sinRemote.sin_port;

  char S[1024];
   if(sprintf(S, "Accepted Connection: %s [Port ID=%d {%d}]", IPFromSocket(Ret), aPort, PortFromSocket(Ret))){
    CSocketSSI::Log(S);
  }
	return Ret;
}

char* CSocketSSI::IPFromSocket(SSI_SOCKET sd)
{
	sockaddr_in an_addr;
	int i = sizeof(an_addr);
	getsockname(sd, (sockaddr*)&an_addr, &i);
	return inet_ntoa(an_addr.sin_addr);
}

int CSocketSSI::PortFromSocket(SSI_SOCKET sd)
{
	sockaddr_in an_addr;
	int i = sizeof(an_addr);
	getsockname(sd, (sockaddr*)&an_addr, &i);
  return an_addr.sin_port;
}

char* CSocketSSI::IPFromName(const char* Name)
{
	struct hostent *pHost = 0;
	struct sockaddr_in SockAddr;
	pHost = gethostbyname(Name);
	memcpy(&SockAddr.sin_addr, pHost->h_addr_list[0], pHost->h_length);
	return inet_ntoa(SockAddr.sin_addr);
}

void CSocketSSI::DispLastError(const char* File, int Line)
{
  int i = CSocketSSI::GetLastError();
	//printf("*** Error from WSAGetLastError: %d ([%d] in %s)\n", i, Line, File);
  char S[1024];
  if(sprintf(S, "*** Error from WSAGetLastError: %d ([%d] in %s)\n", i, Line, File)){
    CSocketSSI::Log(S);
  }
  
	//FILE* OH = fopen("SSI_SOCKET.log", "a");
	//if(OH){
	//	fprintf(OH, "*** Error from WSAGetLastError: %d ([%d] in %s)\n", i, Line, File);
	//	fclose(OH);
	//}
}

int CSocketSSI::GetLastError(){
  int iRet =  WSAGetLastError();
  WSASetLastError(0);
  return iRet;
}


int CSocketSSI::SendData(char* Data, int iLen)
{
	if(iLen < 0){
		return 0;
	}

	if(send(myHandle, Data, iLen, 0) == SOCKET_ERROR){
		DispLastError(__FILE__, __LINE__);
		return 0;
		//Ret.ID = RSP_ERROR;
		//return Ret;
	}

	return 1;
}

int CSocketSSI::ReadData(char* Data, int iLen)
{
	if(iLen < 0){
		return 0;
	}

	int read_bytes = 0;
	int chunk_size = 1;
	chunk_size = iLen;

	while(read_bytes < iLen){
		char* buff = Data;
		int iRead = 0;
    int ReadCtr = 0;
    int bGo = 1;
		//select(
    while(bGo){
		  iRead = recv(myHandle, &buff[read_bytes], chunk_size, 0);
		  if(iRead == SOCKET_ERROR){
			  DispLastError(__FILE__, __LINE__);
			  return 0;
      }else if(iRead == 0){  // Actually, this means connection gracefully closed according to docs
        //printf("Connection Gracefully Closed");
        return 0;
      }else{
        bGo = 0;
      }
    }

		read_bytes += iRead;
		if(read_bytes > iLen){
			return 0;
		}
	}

	return 1;
}

int CSocketSSI::GetHostName(char* Buff, int sz)
{
	return gethostname(Buff, sz);
}


int CSocketSSI::Log(const char* mess){
	//FILE* LogFile = fopen(socket_log_file.c_str(), "a");	
  if(socket_log_file == 0){
    return 0;
  }
  if(strlen(socket_log_file) == 0){
    return 0;
  }
  if(do_log){
	  FILE* LogFile = fopen(socket_log_file, "a");	
	  if(LogFile){
		  fprintf(LogFile, "%s\n", mess);
		  fclose(LogFile);
		  return 1;
	  }else{
		  return 0;
	  }
  }
  return 1;
}

int CSocketSSI::ClearLog(){
  if(do_log){
	  FILE* LogFile = fopen(socket_log_file, "w");	
	  if(LogFile){
		  fclose(LogFile);
	    return 1;
	  }
    return 0;
  }
	return 1;
}

int CSocketSSI::SetLogFile(const char* In){
  if(In == 0){
    return 0;
  }
  int len = strlen(In);
  if(len > 1023){
    return 0;
  }
  if(memcpy(socket_log_file, In, len) == 0){
    return 0;
  }
  do_log = 1;

  char S[512];
  time_t timer;
  time(&timer);
  sprintf(S, "Session: %s", ctime(&timer));
  Log(S);

  return 1;

}